import os
import requests
import datetime
import subprocess
from openai import OpenAI

# Local Qwen3-Coder
# model_name = "kirito1/qwen3-coder"
# url_info = "http://localhost:11434/api/generate"

# Online Deepseek
model_name = "deepseek-reasoner"
url_info = "https://api.deepseek.com/v1"
api_key = "sk-6227423111b14b8087aac2490bcaada3"

def query_ollama(prompt, model=model_name):
    url = url_info
    data = {
        "model": model,
        "prompt": prompt,
        "stream": False
    }
    response = requests.post(url, json=data)
    if response.status_code == 200:
        return response.json()["response"]
    else:
        raise Exception(f"API Error: {response.text}")

def query_online(prompt, model=model_name):
    client = OpenAI(api_key=api_key, base_url=url_info)
    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": "You are a helpful HPC (CUDA CPP) code programmer!"},
            {"role": "user", "content": prompt},
        ],
        stream=False
    )
    return response.choices[0].message.content
    
def create_code_conversion_prompt(original_code, source_lang="Python", target_lang="CUDA-CPP"):
    return f"""Please convert the following {source_lang} code to optimized {target_lang} code with the following requirements:
        1. Include all necessary headers and namespace declarations
        2. Use proper CUDA memory management (cudaMalloc, cudaFree)
        3. Implement error checking for CUDA API calls
        4. Include both host and device code
        5. Add comments explaining key CUDA concepts
        6. Ensure the output is compilable with nvcc
        7. Maintain the exact same functionality as the original code

        {source_lang} Code:
        ```{source_lang.lower()}
        {original_code}```
        
        Please provide only the complete {target_lang} implementation wrapped in a single code block (```). Do not include any explanations outside the code block.
        """

def extract_code_block(response):
    """Extract code block from response with multiple fallback methods"""
    # Case 1: Standard triple backtick code block
    if response.count("```") >= 2:
        code = response.split("```")[1]
        # Remove language specifier if present
        if code.startswith("cuda") or code.startswith("cpp"):
            code = code.split('\n', 1)[1]
        return code
    
    # Case 2: Code wrapped in <code> tags (some models use this)
    if "<code>" in response and "</code>" in response:
        return response.split("<code>")[1].split("</code>")[0]
    
    # Case 3: No delimiters found - try to find the first { and last }
    if "{" in response and "}" in response:
        start = response.find("{")
        end = response.rfind("}")
        return response[start:end+1]
    
    # Final fallback - return the entire response
    return response


def validate_cuda_code(code):
    """Basic validation of CUDA code"""
    required_keywords = [
        "#include", 
        "__global__", 
        "cudaMalloc", 
        "cudaFree",
        "cudaMemcpy"
    ]
    
    missing = [kw for kw in required_keywords if kw not in code]
    if missing:
        raise ValueError(f"Generated code missing required CUDA elements: {missing}")
    
    return code

def sanitize_code(code):
    """Clean up common formatting issues"""
    # Remove leading/trailing whitespace
    code = code.strip()
    
    # Fix common indentation issues
    lines = code.split('\n')
    if len(lines) > 1 and not lines[1].startswith((' ', '\t')):
        # If second line isn't indented, it's likely a formatting issue
        code = '\n'.join([lines[0]] + ['    ' + line for line in lines[1:]])
    
    return code

def generate_and_validate_cuda(python_code, max_retries=3):
    prompt = create_code_conversion_prompt(python_code)
    for attempt in range(max_retries):
        try:
            response = query_online(prompt)
            print("\n- - - - - - - - -\n", response)
            raw_code = extract_code_block(response)
            sanitized_code = sanitize_code(raw_code)
            validated_code = validate_cuda_code(sanitized_code)
            # Additional check for main function
            if "int main(" not in validated_code:
                raise ValueError("Generated code missing main function")
            return validated_code
        except Exception as e:
            print(f"Attempt {attempt + 1} failed: {str(e)}")
            if attempt == max_retries - 1:
                raise
            # Optionally modify prompt based on previous failure
            prompt += f"\n\nPrevious attempt failed because: {str(e)}\nPlease correct these issues."

def generate_cuda_with_feedback(python_code, max_retries=3):
    prompt = create_code_conversion_prompt(python_code)
    previous_errors = []
    for attempt in range(max_retries):
        try:
            if previous_errors:
                error_prompt = "\n\nPrevious errors encountered:\n" + "\n".join(previous_errors)
                prompt += error_prompt + "\nPlease fix these issues in the CUDA code."
            
            response = query_ollama(prompt)
            raw_code = extract_code_block(response)
            sanitized_code = sanitize_code(raw_code)
            validated_code = validate_cuda_code(sanitized_code)
            return validated_code
            
        except Exception as e:
            previous_errors.append(str(e))
            if attempt == max_retries - 1:
                raise
            
def process_file(input_path, output_path):
    with open(input_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    if not content.strip():
        raise ValueError(f"Empty file: {input_path}")
    
    cuda_code = generate_and_validate_cuda(content)
    
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(cuda_code)
    
    return output_path


def compile_and_run_cuda(source_path, executable_path, output_path):
    """
    Return (success: bool, output: str, error: str)
    """
    # 编译命令 - 将错误输出重定向到临时文件
    compile_cmd = f"nvcc -o {executable_path} {source_path} 2>&1"
    
    try:
        # 使用subprocess捕获输出和错误
        compile_process = subprocess.run(
            compile_cmd, 
            shell=True, 
            check=True, 
            text=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT  # 合并stdout和stderr
        )
        
        # 编译成功，运行程序
        run_process = subprocess.run(
            executable_path,
            shell=True,
            text=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        
        # 保存输出
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(run_process.stdout)
            
        return (True, run_process.stdout, "")
        
    except subprocess.CalledProcessError as e:
        # 编译失败，返回错误信息
        error_msg = e.stdout if e.stdout else "Unknown compilation error"
        return (False, "", error_msg)
    except Exception as e:
        # 其他运行错误
        return (False, "", str(e))
    
if __name__ == "__main__":
    dir_path = "./workloads/Sample/"
    base_name = "sliding_quant"
    input_file = f"{dir_path}{base_name}.py"
    cuda_file = f"{dir_path}{base_name}.cu"
    executable = f"{dir_path}{base_name}.exe"
    output_file = f"{dir_path}{base_name}.out"
    error_log = f"{dir_path}compilation_errors.log"
    
    try:
        with open(input_file, 'r', encoding='utf-8') as f:
            python_code = f.read()
        cuda_code = generate_and_validate_cuda(python_code)
        
        with open(cuda_file, 'w', encoding='utf-8') as f:
            f.write(cuda_code)
        
        previous_errors = []
        max_compile_attempts = 4
        
        for attempt in range(max_compile_attempts):
            print(f"\nCompilation attempt Round: {attempt}")
            success, output, error = compile_and_run_cuda(cuda_file, executable, output_file)
            if success:
                print("\n=================\nExecution Output:\n", output)
                break
            print(f"\nCompilation attempt {attempt + 1} failed. Error:\n", error)
            previous_errors.append(error)
            feedback_prompt = f"""
            The previous CUDA code failed to compile with the following error:
            {error}
            Please analyze the error and generate corrected CUDA code.
            Provide only the complete corrected code in a single code block (```).
            """
            response = query_online(feedback_prompt + "\n\nOriginal Python code:\n" + python_code)
            new_code = extract_code_block(response)
            validated_code = validate_cuda_code(sanitize_code(new_code))
            with open(cuda_file, 'w', encoding='utf-8') as f:
                f.write(validated_code)
        else:
            raise RuntimeError(f"Failed after {max_compile_attempts} compilation attempts")
            
    except Exception as e:
        print(f"Error: {str(e)}")
        with open(error_log, 'a') as f:
            f.write(f"{datetime.datetime.now()}: {str(e)}\n")